#include "stdafx.h"
#include "VirtualDisk.h"
#include "filedisk.h"

void FormatLastError(DWORD errorCode, TCHAR* prefix, CString& message)
{
	LPVOID buf;

	FormatMessage(
		FORMAT_MESSAGE_ALLOCATE_BUFFER |
		FORMAT_MESSAGE_FROM_SYSTEM |
		FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL,
		errorCode,
		0,
		(LPTSTR)&buf,
		0,
		NULL
		);
	
	if (prefix)
		message.Append(prefix);	

	message.Append((LPTSTR)buf);

	LocalFree(buf);
}

DWORD GetSystemDrives(CList<CString>& drives)
{
	DWORD  logicalDrives = GetLogicalDrives();
	if (logicalDrives == 0)
		return -1;

	for (DWORD i = 0; i < 26; i++)
	{
		if (logicalDrives & (1 << i))
		{
			CString driveLetter;
			driveLetter.Insert(0, (TCHAR)('A' + i));
			drives.AddTail(driveLetter);
		}
	}

	return 0;
}

DWORD GetFreeDriveLetters(CList<CString>& freeDrives)
{
	DWORD  logicalDrives = GetLogicalDrives();
	if (logicalDrives == 0)
		return -1;

	for (DWORD i = 0; i < 26; i++) 
	{
		if (logicalDrives & (1 << i))
			continue;

		CString driveLetter;
		driveLetter.Insert(0, (TCHAR)('A' + i));
		freeDrives.AddTail(driveLetter);
	}

	return 0;
}

DWORD GetDvdDrives(CList<CString>& dvdDrives)
{
	DWORD  logicalDrives = GetLogicalDrives();
	if (logicalDrives == 0)
		return -1;

	for (DWORD i = 0; i < 26; i++)
	{
		if (!(logicalDrives & (1 << i)))
			continue;

		CString drive;
		drive.Insert(0, (TCHAR)('A' + i));
		drive.Append(_T(":\\"));

		if (GetDriveType(drive.GetString()) != DRIVE_CDROM)
			continue;

		drive.Empty();
		drive.Insert(0, (TCHAR)('A' + i));
		dvdDrives.AddTail(drive);
	}

	return 0;
}

static BOOLEAN isFileDisk(TCHAR driveLetter)
{
	TCHAR volumeName[] = _T("\\\\.\\ :");
	HANDLE device = INVALID_HANDLE_VALUE;
	POPEN_FILE_INFORMATION  fileInfo = NULL;
	DWORD bytesReturned;
	BOOLEAN ret = FALSE;

	do 
	{
		volumeName[4] = driveLetter;

		device = CreateFile(
			volumeName,
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			FILE_FLAG_NO_BUFFERING,
			NULL
			);

		if (device == INVALID_HANDLE_VALUE)
			break;

		fileInfo = (POPEN_FILE_INFORMATION)malloc(sizeof(OPEN_FILE_INFORMATION)+MAX_PATH);
		if (!DeviceIoControl(
			device,
			IOCTL_FILE_DISK_QUERY_FILE,
			NULL,
			0,
			fileInfo,
			sizeof(OPEN_FILE_INFORMATION)+MAX_PATH,
			&bytesReturned,
			NULL
			))
			break;

		if (bytesReturned < sizeof(OPEN_FILE_INFORMATION))
			break;

		ret = TRUE;
	} while (0);

	if (device != INVALID_HANDLE_VALUE)
		CloseHandle(device);

	if (fileInfo)
		free(fileInfo);

	return ret;
}

DWORD GetFileDiskDrives(CList<CString>& fileDiskDrives)
{
	DWORD  logicalDrives = GetLogicalDrives();
	if (logicalDrives == 0)
		return -1;

	for (DWORD i = 0; i < 26; i++)
	{
		if (!(logicalDrives & (1 << i)))
			continue;

		if (!isFileDisk((TCHAR)('A'+i)))
			continue;

		CString drive;
		drive.Insert(0, (TCHAR)('A' + i));
		fileDiskDrives.AddTail(drive);
	}

	return 0;
}

DWORD AddVirtualDrive(int deviceNumber, char* fileName, size_t fileSize, TCHAR driveLetter)
{
	TCHAR  volumeName[] = _T("\\\\.\\ :");
	TCHAR  driveName[] = _T(" :\\");
	TCHAR  deviceName[255];
	HANDLE device = INVALID_HANDLE_VALUE;
	BOOLEAN  undefineDosDrive = FALSE;
	POPEN_FILE_INFORMATION fileInfo = NULL;
	DWORD  bytesReturned;
	DWORD ret = -1;

	volumeName[4] = driveLetter;
	driveName[0] = driveLetter;

	do 
	{
		device = CreateFile(
			volumeName,
			GENERIC_READ | GENERIC_WRITE,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			FILE_FLAG_NO_BUFFERING,
			NULL
			);

		if (device != INVALID_HANDLE_VALUE)
		{
			SetLastError(ERROR_BUSY);
			break;
		}

		_stprintf(deviceName, DEVICE_NAME_PREFIX _T("%u"), deviceNumber);
		if (!DefineDosDevice(
			DDD_RAW_TARGET_PATH,
			&volumeName[4],
			deviceName
			))
			break;

		undefineDosDrive = TRUE;

		device = CreateFile(
			volumeName,
			GENERIC_READ | GENERIC_WRITE,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			FILE_FLAG_NO_BUFFERING,
			NULL
			);

		if (device == INVALID_HANDLE_VALUE)
			break;

		char fileNameBuf[4096] = { 0 };
		_snprintf_c(fileNameBuf, sizeof(fileNameBuf), "\\??\\%s", fileName);
		fileNameBuf[sizeof(fileNameBuf)-1] = 0;
		USHORT fileNameLen = (USHORT)strlen(fileNameBuf);
		size_t fileInfoLen = sizeof(OPEN_FILE_INFORMATION)+fileNameLen - 1;

		fileInfo = (POPEN_FILE_INFORMATION)malloc(fileInfoLen);
		memset(fileInfo, 0, fileInfoLen);

		fileInfo->FileSize.QuadPart = fileSize;
		fileInfo->ReadOnly = FALSE;
		fileInfo->DriveLetter = (UCHAR)driveLetter;
		fileInfo->FileNameLength = fileNameLen;
		memcpy(fileInfo->FileName, fileNameBuf, fileNameLen);

		if (!DeviceIoControl(
			device,
			IOCTL_FILE_DISK_OPEN_FILE,
			fileInfo,
			fileInfoLen,
			NULL,
			0,
			&bytesReturned,
			NULL
			))
			break;

		SHChangeNotify(SHCNE_DRIVEADD, SHCNF_PATH, driveName, NULL);
		 
		undefineDosDrive = FALSE;
		ret = 0;
	} while (0);

	if (device != INVALID_HANDLE_VALUE)
		CloseHandle(device);

	if (fileInfo)
		free(fileInfo);

	if (undefineDosDrive)
		DefineDosDevice(DDD_REMOVE_DEFINITION, &volumeName[4], NULL);
	
	return ret;
}

DWORD RemoveVirtualDrive(TCHAR driveLetter)
{
	TCHAR   volumeName[] = _T("\\\\.\\ :");
	TCHAR   driveName[] = _T(" :\\");
	HANDLE  device = INVALID_HANDLE_VALUE;
	DWORD   bytesReturned;
	DWORD	ret = -1;

	volumeName[4] = driveLetter;
	driveName[0] = driveLetter;

	do 
	{
		device = CreateFile(
			volumeName,
			GENERIC_READ | GENERIC_WRITE,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			FILE_FLAG_NO_BUFFERING,
			NULL
			);

		if (device == INVALID_HANDLE_VALUE)
			break;

		if (!DeviceIoControl(
			device,
			FSCTL_LOCK_VOLUME,
			NULL,
			0,
			NULL,
			0,
			&bytesReturned,
			NULL
			))
			break;

		if (!DeviceIoControl(
			device,
			IOCTL_FILE_DISK_CLOSE_FILE,
			NULL,
			0,
			NULL,
			0,
			&bytesReturned,
			NULL
			))
			break;

		if (!DeviceIoControl(
			device,
			FSCTL_DISMOUNT_VOLUME,
			NULL,
			0,
			NULL,
			0,
			&bytesReturned,
			NULL
			))
			break;

		if (!DeviceIoControl(
			device,
			FSCTL_UNLOCK_VOLUME,
			NULL,
			0,
			NULL,
			0,
			&bytesReturned,
			NULL
			))
			break;

		if (!DefineDosDevice(
			DDD_REMOVE_DEFINITION,
			&volumeName[4],
			NULL
			))
			break;

		SHChangeNotify(SHCNE_DRIVEREMOVED, SHCNF_PATH, driveName, NULL);
		ret = 0;
	} while (0);

	if (device != INVALID_HANDLE_VALUE)
		CloseHandle(device);

	return ret;
}

CString GetDriveLabel(CString selectedDrive)
{
	TCHAR  dvdLabel[MAX_PATH] = { 0 };
	DWORD serialNumber = -1;

	if (!GetVolumeInformation(selectedDrive + ":\\", dvdLabel, MAX_PATH, &serialNumber, NULL, NULL, NULL, 0))
	{
		//CString errorMessage;
		//FormatLastError(GetLastError(), _T("Query CD/DVD drive volume information failed: "), errorMessage);
		// ::MessageBox(NULL, errorMessage, _T("Virtual Disk"), MB_OK);
		return NULL;
	}
	CString result;
	result.Format(TEXT("%s(%d)"), dvdLabel, serialNumber);
	return result;
}